import os
import json
import argparse
import shutil
from utils import ensure_dirs
class ConfigLSTM(object):
    def __init__(self, phase):
        self.is_train = phase == "train"

        self.set_configuration()

        # init hyperparameters and parse from command-line
        parser, args = self.parse()

        # set as attributes
        print("----Experiment Configuration-----")
        for k, v in args.__dict__.items():
            print("{0:20}".format(k), v)
            self.__setattr__(k, v)

        # experiment paths
        self.exp_dir = os.path.join(self.proj_dir, self.exp_name)
        if phase == "train" and args.cont is not True and os.path.exists(self.exp_dir):
            response = input('Experiment log/model already exists, overwrite? (y/n) ')
            if response != 'y':
                exit()
            shutil.rmtree(self.exp_dir)

        self.log_dir = os.path.join(self.exp_dir, 'log')
        self.model_dir = os.path.join(self.exp_dir, 'model')
        ensure_dirs([self.log_dir, self.model_dir])

        # GPU usage
        if args.gpu_ids is not None:
            os.environ["CUDA_VISIBLE_DEVICES"] = str(args.gpu_ids)

        # save this configuration
        if self.is_train:
            with open('{}/config.txt'.format(self.exp_dir), 'w') as f:
                json.dump(args.__dict__, f, indent=2)

    def set_configuration(self):
        self.args_dim = 256  # For LSTM, use your input dimension
        self.n_args = 10  # Example: number of arguments per command, can be changed
        self.n_commands = 50  # Example: number of unique commands

        # LSTM model-specific settings
        self.hidden_dim = 512  # Hidden dimension for LSTM
        self.num_layers = 2    # Number of LSTM layers
        self.dropout = 0.2     # Dropout rate
        self.max_n_loops = 10
        self.max_n_curves = 10
        self.max_total_len = 256

        # Training parameters
        self.lr = 1e-3
        self.nr_epochs = 1000
        self.batch_size = 64
        self.num_workers = 8
        self.grad_clip = 1.0
        self.warmup_step = 2000
        self.cont = False  # Flag to resume training
        self.save_frequency = 500
        self.val_frequency = 10
        self.vis_frequency = 2000
        self.augment = True
        self.use_group_emb = True
        self.d_model = 256
        self.max_num_groups = 30
        self.dim_z = 256

    def parse(self):
        """Initialize argument parser for LSTM configuration"""
        parser = argparse.ArgumentParser()

        parser.add_argument('--proj_dir', type=str, default="proj_log", help="path to project folder where models and logs will be saved")
        parser.add_argument('--data_root', type=str, default="data", help="path to source data folder")
        parser.add_argument('--exp_name', type=str, default=os.getcwd().split('/')[-1], help="name of this experiment")
        parser.add_argument('--gpu_ids', type=str, default='0', help="gpu to use, e.g. 0, 1, 2. CPU not supported.")
        parser.add_argument('--batch_size', type=int, default=self.batch_size, help="batch size")
        parser.add_argument('--num_workers', type=int, default=self.num_workers, help="number of workers for data loading")
        parser.add_argument('--nr_epochs', type=int, default=self.nr_epochs, help="total number of epochs to train")
        parser.add_argument('--lr', type=float, default=self.lr, help="initial learning rate")
        parser.add_argument('--grad_clip', type=float, default=self.grad_clip, help="gradient clipping")
        parser.add_argument('--warmup_step', type=int, default=self.warmup_step, help="step size for learning rate warm up")
        parser.add_argument('--continue', dest='cont', action='store_true', help="continue training from checkpoint")
        parser.add_argument('--ckpt', type=str, default='latest', required=False, help="desired checkpoint to restore")
        parser.add_argument('--vis', action='store_true', default=False, help="visualize output in training")
        parser.add_argument('--save_frequency', type=int, default=self.save_frequency, help="save models every x epochs")
        parser.add_argument('--val_frequency', type=int, default=self.val_frequency, help="run validation every x iterations")
        parser.add_argument('--vis_frequency', type=int, default=self.vis_frequency, help="visualize output every x iterations")
        parser.add_argument('--augment', action='store_true', help="use random data augmentation")

        if not self.is_train:
            parser.add_argument('-m', '--mode', type=str, choices=['rec', 'enc', 'dec'])
            parser.add_argument('-o', '--outputs', type=str, default=None)
            parser.add_argument('--z_path', type=str, default=None)

        args = parser.parse_args()
        return parser, args
